import logging

from openpyxl.reader.excel import load_workbook
from openpyxl.styles import PatternFill

from Util.DBUtil import DBUtil
from sqlalchemy import text
import json
import pandas as pd

querySql = """
SELECT
    id, segment_rules
FROM
    segment 
WHERE
    id IN (
        'SG13390',
        'SG14818',
        'SG14819',
        'SG15139',
        'SG15164',
        'SG15165',
        'SG13389',
        'SG13391',
        'SG13396',
        'SG14816',
        'SG14817',
        'SG13392',
        'SG13393',
        'SG13394',
        'SG13395'
    )
"""


class SegmentService:

    def get_segment_list(self, db_self):
        logging.info(querySql)
        results = db_self.execute_query(text(querySql))
        return results

    def extract_tag_ids_from_segment_rules(self, segment_rules):
        try:
            json_data = json.loads(segment_rules)
        except json.JSONDecodeError as e:
            logging.error("JSON decode error: %s", e)
            return []
        return self.extract_tag_ids(json_data)

    def extract_tag_ids(self, obj):
        tag_ids = []
        if isinstance(obj, dict):
            for key, value in obj.items():
                if key == 'tagId':
                    tag_ids.append(value)
                elif isinstance(value, (dict, list)):
                    tag_ids.extend(self.extract_tag_ids(value))
        elif isinstance(obj, list):
            for item in obj:
                tag_ids.extend(self.extract_tag_ids(item))
        return tag_ids


def write_vertical_tags(data, filename):
    """
    将数据写入 Excel 文件的 'Vertical Tags' 工作表。
    """
    df = pd.DataFrame(data)
    df.to_excel(filename, index=False, sheet_name="Vertical Tags")

    # 加载工作簿以设置样式
    wb = load_workbook(filename)
    ws = wb["Vertical Tags"]

    red_fill = PatternFill(start_color="FF0000", end_color="FF0000", fill_type="solid")

    # 标记缺少 tagId 的行
    for row in ws.iter_rows(min_row=2, max_row=ws.max_row, min_col=1, max_col=1):
        segment_id_cell = row[0]
        tag_id_cell = ws.cell(row=segment_id_cell.row, column=2)
        if tag_id_cell.value is None:
            segment_id_cell.fill = red_fill

    wb.save(filename)
    logging.info(f"'Vertical Tags' sheet updated with missing tagId highlighted.")


def write_horizontal_tags(data, filename):
    """
    将数据写入 Excel 文件的 'Horizontal Tags' 工作表，并标记缺少 tagId 的行。
    """
    segment_ids = [row["Segment ID"] for row in data]
    unique_segment_ids = list(set(segment_ids))
    horizontal_data = []

    for segment_id in unique_segment_ids:
        tags = [row["Tag ID"] for row in data if row["Segment ID"] == segment_id]
        row_data = {"Segment ID": segment_id}
        if any(tag is None for tag in tags):
            row_data["segment_rules"] = [row["segment_rules"] for row in data if row["Segment ID"] == segment_id][0]
        else:
            row_data.update({f"Tag {i + 1}": tag for i, tag in enumerate(tags)})
        horizontal_data.append(row_data)

    df_horizontal = pd.DataFrame(horizontal_data)
    with pd.ExcelWriter(filename, engine='openpyxl', mode='a') as writer:
        df_horizontal.to_excel(writer, index=False, sheet_name="Horizontal Tags")

    # 加载工作簿以设置样式
    wb = load_workbook(filename)
    ws = wb["Horizontal Tags"]

    red_fill = PatternFill(start_color="FF0000", end_color="FF0000", fill_type="solid")

    # 标记缺少 tagId 的行
    for row in ws.iter_rows(min_row=2, max_row=ws.max_row, min_col=1, max_col=1):
        segment_id_cell = row[0]
        if not any(ws.cell(row=segment_id_cell.row, column=col).value for col in range(2, ws.max_column + 1)):
            segment_id_cell.fill = red_fill

    wb.save(filename)
    logging.info(f"'Horizontal Tags' sheet updated with missing tagId highlighted.")

# 运行主函数
if __name__ == "__main__":
    logging.basicConfig(level=logging.INFO)
    # 更改为 "test" 以使用测试环境
    environment = "prod"
    # 更改为适当的数据库名称
    db_name = "audience_multi"
    db_self = DBUtil(environment, db_name);
    segment_service = SegmentService()
    segment_list = segment_service.get_segment_list(db_self)

    data = []
    for row in segment_list:
        segment_id = row[0]
        segment_rules = row[1]
        tag_ids = segment_service.extract_tag_ids_from_segment_rules(segment_rules)
        if tag_ids:
            for tag_id in tag_ids:
                data.append({"Segment ID": segment_id, "Tag ID": tag_id})
        else:
            data.append({"Segment ID": segment_id, "Tag ID": None, "segment_rules": segment_rules})

    output_filename = "output.xlsx"
    write_vertical_tags(data, output_filename)
    write_horizontal_tags(data, output_filename)
